from typing import Optional
from nl2sql.model.llm import LLM
from nl2sql.prompt.query_extract_prompt import query_extract_prompt


class ParseQuery:
    def __init__(self, model: Optional[LLM] = None) -> None:
        """
        提取用户问题中的关键实体
        """
        self.prompt = query_extract_prompt
        self.llm: LLM = model if model is not None else LLM()

    def parse_query(self, query: str) -> Optional[str]:
        """
        用于拆解用户的问题，提取SQL相关实体
        :param query: 用户自然语言问题
        :return: 提取出的实体列表
        """
        filled_prompt = self.prompt.format(query=query)
        llm_resp: str = self.llm.chat(prompt=filled_prompt)
        return llm_resp


if __name__ == "__main__":
    model_name = "qwen-max"
    api_key = "sk-34f5c792513c423f90b404d28b070f1f"
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    llm = LLM(model_name=model_name,
              api_key=api_key,
              base_url=base_url)
    PQ = ParseQuery(llm)
    parser = PQ.parse_query(query="载客输出近一个月的华航，东航，大韩航空，海航这四家中，从江城，山城，魔都起飞的航班载客量")
    print(parser)
